use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use duckdb::Connection;
use ggsql::utils::{sql_generate_series, sql_percentile};

fn setup_connection() -> Connection {
    Connection::open_in_memory().expect("Failed to open DuckDB in-memory connection")
}

fn create_data_table(conn: &Connection, name: &str, n_rows: usize) {
    conn.execute_batch(&format!(
        "CREATE OR REPLACE TABLE {name} AS \
         SELECT random() * 1000.0 AS val \
         FROM GENERATE_SERIES(0, {}) AS seq(n)",
        n_rows - 1
    ))
    .expect("Failed to create data table");
}

fn bench_generate_series(c: &mut Criterion) {
    let conn = setup_connection();
    let mut group = c.benchmark_group("generate_series");

    for n in [64, 512, 1000, 4096] {
        group.bench_with_input(BenchmarkId::new("native", n), &n, |b, &n| {
            let sql = format!(
                "SELECT n FROM GENERATE_SERIES(0, {}) AS seq(n)",
                n - 1
            );
            b.iter(|| {
                let mut stmt = conn.prepare(&sql).unwrap();
                let rows = stmt.query_map([], |row| row.get::<_, f64>(0)).unwrap();
                for r in rows {
                    std::hint::black_box(r.unwrap());
                }
            });
        });

        group.bench_with_input(BenchmarkId::new("portable", n), &n, |b, &n| {
            let cte = sql_generate_series(n);
            let sql = format!("WITH RECURSIVE {cte} SELECT n FROM __ggsql_seq__");
            b.iter(|| {
                let mut stmt = conn.prepare(&sql).unwrap();
                let rows = stmt.query_map([], |row| row.get::<_, f64>(0)).unwrap();
                for r in rows {
                    std::hint::black_box(r.unwrap());
                }
            });
        });
    }

    group.finish();
}

fn bench_percentile(c: &mut Criterion) {
    let conn = setup_connection();
    let mut group = c.benchmark_group("percentile");

    for n_rows in [100, 1000, 10000] {
        let table = format!("data_{n_rows}");
        create_data_table(&conn, &table, n_rows);

        group.bench_with_input(BenchmarkId::new("native", n_rows), &n_rows, |b, _| {
            let sql = format!(
                "SELECT QUANTILE_CONT(val, 0.25) AS q1, QUANTILE_CONT(val, 0.75) AS q3 \
                 FROM {table}"
            );
            b.iter(|| {
                let mut stmt = conn.prepare(&sql).unwrap();
                let row = stmt
                    .query_row([], |row| {
                        Ok((row.get::<_, f64>(0)?, row.get::<_, f64>(1)?))
                    })
                    .unwrap();
                std::hint::black_box(row);
            });
        });

        group.bench_with_input(BenchmarkId::new("portable", n_rows), &n_rows, |b, _| {
            let from = format!("SELECT * FROM {table}");
            let q1 = sql_percentile("val", 0.25, &from, &[]);
            let q3 = sql_percentile("val", 0.75, &from, &[]);
            let sql = format!("SELECT {q1} AS q1, {q3} AS q3");
            b.iter(|| {
                let mut stmt = conn.prepare(&sql).unwrap();
                let row = stmt
                    .query_row([], |row| {
                        Ok((row.get::<_, f64>(0)?, row.get::<_, f64>(1)?))
                    })
                    .unwrap();
                std::hint::black_box(row);
            });
        });
    }

    group.finish();
}

criterion_group!(benches, bench_generate_series, bench_percentile);
criterion_main!(benches);